#include "nnie.h"

#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <pthread.h>
#include "hi_comm_svp.h"
#include "sample_comm_svp.h"
#include "sample_comm_ive.h"
#include "sample_comm_nnie.h"
#include "sample_svp_nnie_software.h"

#include "defines.h"
#include "mpp.h"
#include "util.h"

#define MODEL_FILE      "./data/yolov3_trash_detection_inst.wk"
#define IMAGE_WIDTH     416
#define IMAGE_HEIGHT    416
#define USLEEP_TIME     100
#define DETECT_OBJ_MAX  64
#define CONF_THRESHOLD  0.40f
#define SCORE_THRESHOLD 0.25f

static const char *LABEL_NAMES[] = {
    "background", // 为啥还有背景呢
    "一次性快餐盒", "书籍纸张", "充电宝", "剩饭剩菜", "包", "垃圾桶", "塑料器皿", "塑料玩具", "塑料衣架", "大骨头", "干电池",
    "快递纸袋", "插头电线", "旧衣服", "易拉罐", "枕头", "果皮果肉", "毛绒玩具", "污损塑料", "污损用纸", "洗护用品", "烟蒂",
    "牙签", "玻璃器皿", "砧板", "筷子", "纸盒纸箱", "花盆", "茶叶渣", "菜帮菜叶", "蛋壳", "调料瓶", "软膏", "过期药物",
    "酒瓶", "金属厨具", "金属器皿", "金属食品罐", "锅", "陶瓷器皿", "鞋", "食用油桶", "饮料瓶", "鱼骨"
};

static SAMPLE_SVP_NNIE_MODEL_S g_stNnieModel;
static SAMPLE_SVP_NNIE_PARAM_S g_stNnieParam;
static SAMPLE_SVP_NNIE_YOLOV3_SOFTWARE_PARAM_S g_stYolov3Param;
static int g_bIsRunning;
static pthread_t g_pidAiThread;
static int g_bEnableForward;

static HI_S32 SVP_NNIE_YOLOV3_ParamInit(SAMPLE_SVP_NNIE_PARAM_S *pstNnieParam, SAMPLE_SVP_NNIE_YOLOV3_SOFTWARE_PARAM_S* pstYolov3Param) {
    HI_S32 s32Ret = 0;

    // 初始化yolov3参数
    pstYolov3Param->u32OriImHeight = pstNnieParam->astSegData[0].astSrc[0].unShape.stWhc.u32Height;
    pstYolov3Param->u32OriImWidth = pstNnieParam->astSegData[0].astSrc[0].unShape.stWhc.u32Width;
    pstYolov3Param->u32BboxNumEachGrid = 3;
    pstYolov3Param->u32ClassNum = 44;
    // sample_svp_nnie_software.h 中的 SAMPLE_SVP_NNIE_YOLOV3_EACH_BBOX_INFER_RESULT_NUM 需要改成 u32ClassNum+5
    pstYolov3Param->au32GridNumHeight[0] = 13;
    pstYolov3Param->au32GridNumHeight[1] = 26;
    pstYolov3Param->au32GridNumHeight[2] = 52;
    pstYolov3Param->au32GridNumWidth[0] = 13;
    pstYolov3Param->au32GridNumWidth[1] = 26;
    pstYolov3Param->au32GridNumWidth[2] = 52;
    pstYolov3Param->u32NmsThresh = (HI_U32) (0.3f * SAMPLE_SVP_QUANT_BASE);
    pstYolov3Param->u32ConfThresh = (HI_U32) (0.25f * SAMPLE_SVP_QUANT_BASE);
    pstYolov3Param->u32MaxRoiNum = 10;
    pstYolov3Param->af32Bias[0][0] = 116;
    pstYolov3Param->af32Bias[0][1] = 90;
    pstYolov3Param->af32Bias[0][2] = 156;
    pstYolov3Param->af32Bias[0][3] = 198;
    pstYolov3Param->af32Bias[0][4] = 373;
    pstYolov3Param->af32Bias[0][5] = 326;
    pstYolov3Param->af32Bias[1][0] = 30;
    pstYolov3Param->af32Bias[1][1] = 61;
    pstYolov3Param->af32Bias[1][2] = 62;
    pstYolov3Param->af32Bias[1][3] = 45;
    pstYolov3Param->af32Bias[1][4] = 59;
    pstYolov3Param->af32Bias[1][5] = 119;
    pstYolov3Param->af32Bias[2][0] = 10;
    pstYolov3Param->af32Bias[2][1] = 13;
    pstYolov3Param->af32Bias[2][2] = 16;
    pstYolov3Param->af32Bias[2][3] = 30;
    pstYolov3Param->af32Bias[2][4] = 33;
    pstYolov3Param->af32Bias[2][5] = 23;

    // 申请辅助内存空间
    HI_U64 u64PhyAddr = 0;
    HI_U8* pu8VirAddr = NULL;
    HI_U32 u32ClassNum = pstYolov3Param->u32ClassNum + 1;
    HI_U32 u32TmpBufTotalSize = SAMPLE_SVP_NNIE_Yolov3_GetResultTmpBuf(pstNnieParam, pstYolov3Param);
    HI_U32 u32DstRoiSize = SAMPLE_SVP_NNIE_ALIGN16(u32ClassNum * pstYolov3Param->u32MaxRoiNum * sizeof(HI_U32) * SAMPLE_SVP_COORDI_NUM);
    HI_U32 u32DstScoreSize = SAMPLE_SVP_NNIE_ALIGN16(u32ClassNum * pstYolov3Param->u32MaxRoiNum * sizeof(HI_U32));
    HI_U32 u32ClassRoiNumSize = SAMPLE_SVP_NNIE_ALIGN16(u32ClassNum * sizeof(HI_U32));
    HI_U32 u32TotalSize = u32TmpBufTotalSize + u32DstRoiSize + u32DstScoreSize + u32ClassRoiNumSize;
    s32Ret = SAMPLE_COMM_SVP_MallocCached("SAMPLE_YOLOV3_INIT", NULL, (HI_U64*) &u64PhyAddr, (void**) &pu8VirAddr, u32TotalSize);
    if (s32Ret != HI_SUCCESS) {
        LOG("Error: alloc memory failed!\n");
        return s32Ret;
    }
    memset_s(pu8VirAddr, u32TotalSize, 0, u32TotalSize);
    SAMPLE_COMM_SVP_FlushCache(u64PhyAddr, (void*) pu8VirAddr, u32TotalSize);
    pstYolov3Param->stGetResultTmpBuf.u64PhyAddr = u64PhyAddr;
    pstYolov3Param->stGetResultTmpBuf.u64VirAddr = SAMPLE_SVP_NNIE_CONVERT_PTR_TO_ADDR(HI_U64, pu8VirAddr);

    // 设置结果blob
    pstYolov3Param->stDstRoi.enType = SVP_BLOB_TYPE_S32;
    pstYolov3Param->stDstRoi.u64PhyAddr = u64PhyAddr + u32TmpBufTotalSize;
    pstYolov3Param->stDstRoi.u64VirAddr = SAMPLE_SVP_NNIE_CONVERT_PTR_TO_ADDR(HI_U64, pu8VirAddr + u32TmpBufTotalSize);
    pstYolov3Param->stDstRoi.u32Stride = SAMPLE_SVP_NNIE_ALIGN16(u32ClassNum * pstYolov3Param->u32MaxRoiNum * sizeof(HI_U32) * SAMPLE_SVP_COORDI_NUM);
    pstYolov3Param->stDstRoi.u32Num = 1;
    pstYolov3Param->stDstRoi.unShape.stWhc.u32Chn = 1;
    pstYolov3Param->stDstRoi.unShape.stWhc.u32Height = 1;
    pstYolov3Param->stDstRoi.unShape.stWhc.u32Width = u32ClassNum * pstYolov3Param->u32MaxRoiNum * SAMPLE_SVP_COORDI_NUM;

    pstYolov3Param->stDstScore.enType = SVP_BLOB_TYPE_S32;
    pstYolov3Param->stDstScore.u64PhyAddr = u64PhyAddr + u32TmpBufTotalSize + u32DstRoiSize;
    pstYolov3Param->stDstScore.u64VirAddr = SAMPLE_SVP_NNIE_CONVERT_PTR_TO_ADDR(HI_U64, pu8VirAddr + u32TmpBufTotalSize + u32DstRoiSize);
    pstYolov3Param->stDstScore.u32Stride = SAMPLE_SVP_NNIE_ALIGN16(u32ClassNum * pstYolov3Param->u32MaxRoiNum * sizeof(HI_U32));
    pstYolov3Param->stDstScore.u32Num = 1;
    pstYolov3Param->stDstScore.unShape.stWhc.u32Chn = 1;
    pstYolov3Param->stDstScore.unShape.stWhc.u32Height = 1;
    pstYolov3Param->stDstScore.unShape.stWhc.u32Width = u32ClassNum * pstYolov3Param->u32MaxRoiNum;

    pstYolov3Param->stClassRoiNum.enType = SVP_BLOB_TYPE_S32;
    pstYolov3Param->stClassRoiNum.u64PhyAddr = u64PhyAddr + u32TmpBufTotalSize + u32DstRoiSize + u32DstScoreSize;
    pstYolov3Param->stClassRoiNum.u64VirAddr = SAMPLE_SVP_NNIE_CONVERT_PTR_TO_ADDR(HI_U64, pu8VirAddr + u32TmpBufTotalSize + u32DstRoiSize + u32DstScoreSize);
    pstYolov3Param->stClassRoiNum.u32Stride = SAMPLE_SVP_NNIE_ALIGN16(u32ClassNum * sizeof(HI_U32));
    pstYolov3Param->stClassRoiNum.u32Num = 1;
    pstYolov3Param->stClassRoiNum.unShape.stWhc.u32Chn = 1;
    pstYolov3Param->stClassRoiNum.unShape.stWhc.u32Height = 1;
    pstYolov3Param->stClassRoiNum.unShape.stWhc.u32Width = u32ClassNum;

    return s32Ret;
}

static HI_S32 SVP_NNIE_FillSrcData(SAMPLE_SVP_NNIE_PARAM_S *pstNnieParam,
                                    SAMPLE_SVP_NNIE_INPUT_DATA_INDEX_S *pstInputDataIdx,
                                    const IVE_IMAGE_S *img) {
    HI_U32 segId = pstInputDataIdx->u32SegIdx;
    HI_U32 nodeId = pstInputDataIdx->u32NodeIdx;

    // 获取数据大小
    HI_U32 u32VarSize;
    if (pstNnieParam->astSegData[segId].astSrc[nodeId].enType >= SVP_BLOB_TYPE_U8 &&
        pstNnieParam->astSegData[segId].astSrc[nodeId].enType <= SVP_BLOB_TYPE_YVU422SP) {
        u32VarSize = sizeof(HI_U8);
    } else {
        u32VarSize = sizeof(HI_U32);
    }

    // 填充源数据
    if (pstNnieParam->astSegData[segId].astSrc[nodeId].enType == SVP_BLOB_TYPE_YVU422SP ||
        pstNnieParam->astSegData[segId].astSrc[nodeId].enType == SVP_BLOB_TYPE_SEQ_S32) {
        LOG("Error: not support this blob type: %d!\n", pstNnieParam->astSegData[segId].astSrc[nodeId].enType);
        return HI_FAILURE;
    }
    if (pstNnieParam->astSegData[segId].astSrc[nodeId].u32Num != 1) {
        LOG("Error: batch size must be 1!\n");
        return HI_FAILURE;
    }

    HI_U32 u32Height = pstNnieParam->astSegData[segId].astSrc[nodeId].unShape.stWhc.u32Height;
    HI_U32 u32Width = pstNnieParam->astSegData[segId].astSrc[nodeId].unShape.stWhc.u32Width;
    HI_U32 u32Chn = pstNnieParam->astSegData[segId].astSrc[nodeId].unShape.stWhc.u32Chn;
    HI_U32 u32Stride = pstNnieParam->astSegData[segId].astSrc[nodeId].u32Stride;
    HI_U8 *pu8PicAddr = SAMPLE_SVP_NNIE_CONVERT_64BIT_ADDR(HI_U8, pstNnieParam->astSegData[segId].astSrc[nodeId].u64VirAddr);

    if (pstNnieParam->astSegData[segId].astSrc[nodeId].enType == SVP_BLOB_TYPE_YVU420SP) {
        const HI_U8 *srcData = NULL;
        // Y
        srcData = SAMPLE_SVP_NNIE_CONVERT_64BIT_ADDR(const HI_U8, img->au64VirAddr[0]);
        for (HI_U32 j = 0; j < u32Height; j++) {
            if (memcpy_s(pu8PicAddr, u32Width * u32VarSize, srcData, u32Width * u32VarSize) != EOK) {
                LOG("Error: memcpy failed!\n");
                return HI_FAILURE;
            }
            pu8PicAddr += u32Stride;
            srcData += img->au32Stride[0];
        }
        // UV
        srcData = SAMPLE_SVP_NNIE_CONVERT_64BIT_ADDR(const HI_U8, img->au64VirAddr[1]);
        for (HI_U32 j = 0; j < u32Height / 2; j++) { // 2: 1/2Height
            if (memcpy_s(pu8PicAddr, u32Width * u32VarSize, srcData, u32Width * u32VarSize) != EOK) {
                LOG("Error: memcpy failed!\n");
                return HI_FAILURE;
            }
            pu8PicAddr += u32Stride;
            srcData += img->au32Stride[1];
        }
    } else {
        for (HI_U32 i = 0; i < u32Chn; i++) {
            const HI_U8 *srcData = SAMPLE_SVP_NNIE_CONVERT_64BIT_ADDR(const HI_U8, img->au64VirAddr[i]);
            for (HI_U32 j = 0; j < u32Height; j++) {
                if (memcpy_s(pu8PicAddr, u32Width * u32VarSize, srcData, u32Width * u32VarSize) != EOK) {
                    LOG("Error: memcpy failed!\n");
                    return HI_FAILURE;
                }
                pu8PicAddr += u32Stride;
                srcData += img->au32Stride[i];
            }
        }
    }

    // 刷新cache
    HI_S32 s32Ret = SAMPLE_COMM_SVP_FlushCache(pstNnieParam->astSegData[segId].astSrc[nodeId].u64PhyAddr,
        SAMPLE_SVP_NNIE_CONVERT_64BIT_ADDR(HI_VOID, pstNnieParam->astSegData[segId].astSrc[nodeId].u64VirAddr),
        pstNnieParam->astSegData[segId].astSrc[nodeId].u32Num * u32Chn * u32Height * u32Stride);
    if (s32Ret != HI_SUCCESS) {
        LOG("SAMPLE_COMM_SVP_FlushCache failed with %#x!\n", s32Ret);
        return HI_FAILURE;
    }

    return HI_SUCCESS;
}

static HI_S32 SVP_NNIE_Forward(SAMPLE_SVP_NNIE_PARAM_S *pstNnieParam,
                                SAMPLE_SVP_NNIE_INPUT_DATA_INDEX_S* pstInputDataIdx,
                                SAMPLE_SVP_NNIE_PROCESS_SEG_INDEX_S* pstProcSegIdx,
                                HI_BOOL bInstant) {
    HI_S32 s32Ret = 0;
    HI_BOOL bFinish = HI_FALSE;
    SVP_NNIE_HANDLE hSvpNnieHandle = 0;
    HI_U32 u32TotalStepNum = 0;

    SAMPLE_COMM_SVP_FlushCache(pstNnieParam->astForwardCtrl[pstProcSegIdx->u32SegIdx].stTskBuf.u64PhyAddr,
        SAMPLE_SVP_NNIE_CONVERT_64BIT_ADDR(HI_VOID, pstNnieParam->astForwardCtrl[pstProcSegIdx->u32SegIdx].stTskBuf.u64VirAddr),
        pstNnieParam->astForwardCtrl[pstProcSegIdx->u32SegIdx].stTskBuf.u32Size);

    for (HI_U32 i = 0; i < pstNnieParam->astForwardCtrl[pstProcSegIdx->u32SegIdx].u32DstNum; i++) {
        if (pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst[i].enType == SVP_BLOB_TYPE_SEQ_S32) {
            for (HI_U32 j = 0; j < pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst[i].u32Num; j++) {
                u32TotalStepNum += *(SAMPLE_SVP_NNIE_CONVERT_64BIT_ADDR(HI_U32,
                    pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst[i].unShape.stSeq.u64VirAddrStep) + j);
            }
            SAMPLE_COMM_SVP_FlushCache(pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst[i].u64PhyAddr,
                SAMPLE_SVP_NNIE_CONVERT_64BIT_ADDR(HI_VOID, pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst[i].u64VirAddr),
                u32TotalStepNum * pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst[i].u32Stride);
        } else {
            SAMPLE_COMM_SVP_FlushCache(pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst[i].u64PhyAddr,
                SAMPLE_SVP_NNIE_CONVERT_64BIT_ADDR(HI_VOID, pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst[i].u64VirAddr),
                pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst[i].u32Num *
                pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst[i].unShape.stWhc.u32Chn *
                pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst[i].unShape.stWhc.u32Height *
                pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst[i].u32Stride);
        }
    }

    // 根据节点名称设置输入blob
    if (pstInputDataIdx->u32SegIdx != pstProcSegIdx->u32SegIdx) {
        for (HI_U32 i = 0; i < pstNnieParam->pstModel->astSeg[pstProcSegIdx->u32SegIdx].u16SrcNum; i++) {
            HI_U32 j;
            for (j = 0; j < pstNnieParam->pstModel->astSeg[pstInputDataIdx->u32SegIdx].u16DstNum; j++) {
                if (strncmp(pstNnieParam->pstModel->astSeg[pstInputDataIdx->u32SegIdx].astDstNode[j].szName,
                    pstNnieParam->pstModel->astSeg[pstProcSegIdx->u32SegIdx].astSrcNode[i].szName, SVP_NNIE_NODE_NAME_LEN) == 0) {
                    pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astSrc[i] = pstNnieParam->astSegData[pstInputDataIdx->u32SegIdx].astDst[j];
                    break;
                }
            }
            if (j == pstNnieParam->pstModel->astSeg[pstInputDataIdx->u32SegIdx].u16DstNum) {
                LOG("Error: can't find %d-th seg's %d-th src blob!\n", pstProcSegIdx->u32SegIdx, i);
                return HI_FAILURE;
            }
        }
    }

    // 多节点输入输出的CNN类型网络预测
    s32Ret = HI_MPI_SVP_NNIE_Forward(&hSvpNnieHandle,
        pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astSrc,
        pstNnieParam->pstModel, pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst,
        &pstNnieParam->astForwardCtrl[pstProcSegIdx->u32SegIdx], bInstant);
    if (s32Ret != HI_SUCCESS) {
        LOG("Error: HI_MPI_SVP_NNIE_Forward failed with %#x!\n", s32Ret);
        return s32Ret;
    }

    if (bInstant) {
        // 查询任务是否完成
        while ((s32Ret = HI_MPI_SVP_NNIE_Query(pstNnieParam->astForwardCtrl[pstProcSegIdx->u32SegIdx].enNnieId,
                    hSvpNnieHandle, &bFinish, HI_TRUE)) == HI_ERR_SVP_NNIE_QUERY_TIMEOUT) {
            usleep(USLEEP_TIME);
            // LOG("HI_MPI_SVP_NNIE_Query Query timeout!\n");
        }
    }

    u32TotalStepNum = 0;
    for (HI_U32 i = 0; i < pstNnieParam->astForwardCtrl[pstProcSegIdx->u32SegIdx].u32DstNum; i++) {
        if (pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst[i].enType == SVP_BLOB_TYPE_SEQ_S32) {
            for (HI_U32 j = 0; j < pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst[i].u32Num; j++) {
                u32TotalStepNum += *(SAMPLE_SVP_NNIE_CONVERT_64BIT_ADDR(HI_U32,
                    pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst[i].unShape.stSeq.u64VirAddrStep) + j);
            }
            SAMPLE_COMM_SVP_FlushCache(pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst[i].u64PhyAddr,
                SAMPLE_SVP_NNIE_CONVERT_64BIT_ADDR(HI_VOID, pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst[i].u64VirAddr),
                u32TotalStepNum*pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst[i].u32Stride);
        } else {
            SAMPLE_COMM_SVP_FlushCache(pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst[i].u64PhyAddr,
                SAMPLE_SVP_NNIE_CONVERT_64BIT_ADDR(HI_VOID, pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst[i].u64VirAddr),
                pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst[i].u32Num *
                pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst[i].unShape.stWhc.u32Chn *
                pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst[i].unShape.stWhc.u32Height *
                pstNnieParam->astSegData[pstProcSegIdx->u32SegIdx].astDst[i].u32Stride);
        }
    }

    return s32Ret;
}

static void SVP_NNIE_YOLOV3_ParseResult(SVP_BLOB_S *pstDstScore, SVP_BLOB_S *pstDstRoi, SVP_BLOB_S *pstClassRoiNum,
    DetectObjInfo resBuf[], HI_S32 resMaxSize, HI_S32* resLen) {
    HI_U32 u32RoiNumBias = 0;
    HI_U32 u32ScoreBias = 0;
    HI_U32 u32BboxBias = 0;
    HI_FLOAT f32Score = 0.0f;
    HI_S32* ps32Score = SAMPLE_SVP_NNIE_CONVERT_64BIT_ADDR(HI_S32, pstDstScore->u64VirAddr);
    HI_S32* ps32Roi = SAMPLE_SVP_NNIE_CONVERT_64BIT_ADDR(HI_S32, pstDstRoi->u64VirAddr);
    HI_S32* ps32ClassRoiNum = SAMPLE_SVP_NNIE_CONVERT_64BIT_ADDR(HI_S32, pstClassRoiNum->u64VirAddr);
    HI_U32 u32ClassNum = pstClassRoiNum->unShape.stWhc.u32Width;

    HI_S32 resId = 0;
    memset_s(resBuf, resMaxSize * sizeof(resBuf[0]), 0, resMaxSize * sizeof(resBuf[0]));
    *resLen = 0;

    u32RoiNumBias += ps32ClassRoiNum[0];
    for (HI_U32 i = 1; i < u32ClassNum; i++) {
        u32ScoreBias = u32RoiNumBias;
        u32BboxBias = u32RoiNumBias * SAMPLE_SVP_COORDI_NUM;
        // 如果置信度分数大于结果阈值，则打印结果
        if ((HI_FLOAT) ps32Score[u32ScoreBias] / SAMPLE_SVP_QUANT_BASE >= CONF_THRESHOLD && ps32ClassRoiNum[i] != 0) {
            for (HI_U32 j = 0; j < (HI_U32) ps32ClassRoiNum[i]; j++) {
                f32Score = (HI_FLOAT) ps32Score[u32ScoreBias + j] / SAMPLE_SVP_QUANT_BASE;
                    if (f32Score >= SCORE_THRESHOLD) {
                    if (resId >= resMaxSize) {
                        LOG("Warning: resBuf is full!\n");
                        break;
                    }
                    resBuf[resId].cls = i;
                    resBuf[resId].score = f32Score;
                    RectBox *box = &resBuf[resId].box;
                    box->xmin = ps32Roi[u32BboxBias + j * SAMPLE_SVP_COORDI_NUM];
                    box->ymin = ps32Roi[u32BboxBias + j * SAMPLE_SVP_COORDI_NUM + 1];
                    box->xmax = ps32Roi[u32BboxBias + j * SAMPLE_SVP_COORDI_NUM + 2];
                    box->ymax = ps32Roi[u32BboxBias + j * SAMPLE_SVP_COORDI_NUM + 3];
                    if (box->xmin < box->xmax && box->ymin < box->ymax) {
                        resId++;
                    } else {
                        LOG("object %d: %f, (%d, %d), (%d, %d) discard for coord ERR\n",
                            i, f32Score, box->xmin, box->ymin, box->xmax, box->ymax);
                    }
                }
            }
            u32RoiNumBias += ps32ClassRoiNum[i];
        }
    }

    *resLen = resId;
}

static HI_S32 SVP_NNIE_YOLOV3_Forward(const IVE_IMAGE_S *img, DetectObjInfo *resBuf, int resMaxSize, int *resLen) {
    HI_S32 s32Ret = 0;
    SAMPLE_SVP_NNIE_INPUT_DATA_INDEX_S stInputDataIdx = {0};
    SAMPLE_SVP_NNIE_PROCESS_SEG_INDEX_S stProcSegIdx = {0};

    // 填充源数据
    stInputDataIdx.u32SegIdx = 0;
    stInputDataIdx.u32NodeIdx = 0;
    s32Ret = SVP_NNIE_FillSrcData(&g_stNnieParam, &stInputDataIdx, img);
    if (s32Ret != HI_SUCCESS) {
        LOG("SVP_NNIE_FillSrcData failed with %#x!\n", s32Ret);
        return s32Ret;
    }

    // NNIE推理(process the 0-th segment)
    stProcSegIdx.u32SegIdx = 0;
    s32Ret = SVP_NNIE_Forward(&g_stNnieParam, &stInputDataIdx, &stProcSegIdx, HI_TRUE);
    if (s32Ret != HI_SUCCESS) {
        LOG("SVP_NNIE_Forward failed with %#x!\n", s32Ret);
        return s32Ret;
    }

    // CPU后处理
    s32Ret = SAMPLE_SVP_NNIE_Yolov3_GetResult(&g_stNnieParam, &g_stYolov3Param);
    if (s32Ret != HI_SUCCESS) {
        LOG("SAMPLE_SVP_NNIE_Yolov3_GetResult failed with %#x!\n", s32Ret);
        return s32Ret;
    }
    SVP_NNIE_YOLOV3_ParseResult(&g_stYolov3Param.stDstScore, &g_stYolov3Param.stDstRoi, &g_stYolov3Param.stClassRoiNum,
                                    resBuf, resMaxSize, resLen);

    return HI_SUCCESS;
}

static void* nnie_process(void* arg) {
    HI_S32 s32Ret = 0;
    VIDEO_FRAME_INFO_S stFrame = {0};
    VIDEO_FRAME_INFO_S stResizedFrame = {0};
    IVE_IMAGE_S stImage = {0};
    DetectObjInfo objs[DETECT_OBJ_MAX] = {0};
    RectBox objBoxs[DETECT_OBJ_MAX] = {0};
    HI_S32 objNum = 0;

    while (g_bIsRunning) {
        s32Ret = HI_MPI_VPSS_GetChnFrame(0, 1, &stFrame, 2000);
        if (s32Ret != HI_SUCCESS) {
            LOG("HI_MPI_VPSS_GetChnFrame failed with %#x!\n", s32Ret);
            goto release;
        }

        if (g_bEnableForward) {
            s32Ret = MppFrmResize(&stFrame, &stResizedFrame, IMAGE_WIDTH, IMAGE_HEIGHT);
            if (s32Ret != HI_SUCCESS) {
                LOG("MppFrmResize failed with %#x!\n", s32Ret);
                goto destroy_resized_frm;
            }

            s32Ret = FrmToOrigImg(&stResizedFrame, &stImage);
            if (s32Ret != HI_SUCCESS) {
                LOG("FrmToOrigImg failed with %#x!\n", s32Ret);
                goto destroy_resized_frm;
            }

            s32Ret = SVP_NNIE_YOLOV3_Forward(&stImage, objs, DETECT_OBJ_MAX, &objNum);
            if (s32Ret != HI_SUCCESS) {
                LOG("SVP_NNIE_YOLOV3_Forward failed with %#x!\n", s32Ret);
                goto destroy_resized_frm;
            }

            // LOG("YOLOv3 detection result:\n");
            for (HI_S32 i = 0; i < objNum; i++) {
                RectBox *box = &objs[i].box;
                RectBoxTran(box, stResizedFrame.stVFrame.u32Width, stResizedFrame.stVFrame.u32Height,
                    stFrame.stVFrame.u32Width, stFrame.stVFrame.u32Height);
                objBoxs[i] = *box;
                // LOG("%s: %f, (%d, %d), (%d, %d)\n", LABEL_NAMES[objs[i].cls], objs[i].score,
                //         box->xmin, box->ymin, box->xmax, box->ymax);
            }
            nnie_callback(objs, objNum);

            if (objNum > 0) {
                MppFrmDrawRects(&stFrame, objBoxs, objNum, RGB888_RED, 4);
            }
destroy_resized_frm:
            MppFrmDestroy(&stResizedFrame);
        }

send:
        s32Ret = HI_MPI_VO_SendFrame(0, 0, &stFrame, 0);
        if (s32Ret != HI_SUCCESS) {
            LOG("HI_MPI_VO_SendFrame failed with %#x!\n", s32Ret);
            goto release;
        }
release:
        s32Ret = HI_MPI_VPSS_ReleaseChnFrame(0, 1, &stFrame);
        if (s32Ret != HI_SUCCESS) {
            LOG("HI_MPI_VPSS_ReleaseChnFrame failed with %#x!\n", s32Ret);
        }
    }

    return NULL;
}

int nnie_init() {
    HI_S32 s32Ret = 0;
    SAMPLE_SVP_NNIE_CFG_S stNnieCfg = {
        .pszPic = NULL,
        .u32MaxInputNum = 1,
        .u32MaxRoiNum = 0,
        .aenNnieCoreId[0] = SVP_NNIE_ID_0
    };

    // s32Ret = SAMPLE_COMM_SVP_CheckSysInit(); // 不需要, SYS已经初始化过了
    // if (s32Ret != HI_SUCCESS) {
    //     goto init_sys_failed;
    // }
    s32Ret = SAMPLE_COMM_SVP_NNIE_LoadModel(MODEL_FILE, &g_stNnieModel);
    if (s32Ret != HI_SUCCESS) {
        LOG("Error: SAMPLE_COMM_SVP_NNIE_LoadModel failed!\n");
        goto load_model_failed;
    }
    g_stNnieParam.pstModel = &g_stNnieModel.stModel;
    s32Ret = SAMPLE_COMM_SVP_NNIE_ParamInit(&stNnieCfg, &g_stNnieParam);
    if (s32Ret != HI_SUCCESS) {
        LOG("Error: SAMPLE_COMM_SVP_NNIE_ParamInit failed!\n");
        goto init_nnie_failed;
    }
    s32Ret = SVP_NNIE_YOLOV3_ParamInit(&g_stNnieParam, &g_stYolov3Param);
    if (s32Ret != HI_SUCCESS) {
        LOG("Error: SVP_NNIE_YOLOV3_ParamInit failed!\n");
        goto init_yolov3_failed;
    }

    g_bIsRunning = 1;
    s32Ret = pthread_create(&g_pidAiThread, NULL, nnie_process, NULL);
    if (s32Ret != 0) {
        LOG("Error: create nnie thread failed!\n");
        goto create_thread_failed;
    }

    return s32Ret;

create_thread_failed:
    g_bIsRunning = 0;
init_yolov3_failed:
    SAMPLE_SVP_MMZ_FREE(g_stYolov3Param.stGetResultTmpBuf.u64PhyAddr, g_stYolov3Param.stGetResultTmpBuf.u64VirAddr);    
init_nnie_failed:
    SAMPLE_COMM_SVP_NNIE_ParamDeinit(&g_stNnieParam);
load_model_failed:
    SAMPLE_COMM_SVP_NNIE_UnloadModel(&g_stNnieModel);
init_sys_failed:
    return s32Ret;
}

void nnie_exit() {
    g_bIsRunning = 0;
    pthread_join(g_pidAiThread, NULL);
    SAMPLE_SVP_MMZ_FREE(g_stYolov3Param.stGetResultTmpBuf.u64PhyAddr, g_stYolov3Param.stGetResultTmpBuf.u64VirAddr);
    SAMPLE_COMM_SVP_NNIE_ParamDeinit(&g_stNnieParam);
    SAMPLE_COMM_SVP_NNIE_UnloadModel(&g_stNnieModel);
}

int nnie_start() {
    if (!g_bEnableForward) {
        g_bEnableForward = 1;
        return SUCCESS;
    }
    return FAILED;
}

int nnie_stop() {
    if (g_bEnableForward) {
        g_bEnableForward = 0;
        return SUCCESS;
    }
    return FAILED;
}
